"""
 注意：：：这个版本 测试如果不使用open ai 的大模型 返回的格式不会遵循json（应该是.with_structured_output(xxx)不支持） ，导致运行出错。所以另外基于这个改了一下提示词版本的。

reAct 大模型的一种思维模式，让大模型先思考再行动。（Reason + Act） 框架。
主要工具：
    tavily 搜索检索 ，需要配置api
    langsmit  agent 监控工具 ，需要配置api

本示例：
    使用 LangGraph 构建一个完整的“ReAct智能体工作流”。
    - Tavily 搜索工具：让模型具备网络搜索能力。
    - LangSmith（可选）：用于可视化 agent 的执行过程。

"""
import asyncio
import operator
from typing import TypedDict, List, Annotated, Tuple, Union, Literal

from dotenv import load_dotenv
from langchain import hub
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from pydantic import Field, BaseModel

# 自动加载 .env 文件中的环境变量
load_dotenv()

# 创建 Tavily 搜索工具（可联网搜索），每次只返回一个搜索结果
tools = [TavilySearchResults(max_results=1)]

# 从 LangChain Hub 获取 ReAct 框架提示模板
# Hub 上保存了一些通用的 prompt，可以直接 pull 使用
prompt_template = """
你是一个专业 Agent，拥有强大的搜索能力。
你的目标是执行用户提供的任务。
你必须严格遵循 ReAct 模式（Thought, Action, Observation）。

你拥有的工具是：{tool_names}

请使用以下格式进行响应：

Thought: 我应该思考如何执行任务。
Action: tool_name
Action Input: tool_input

... (重复 Thought/Action/Observation)

Thought: 我知道最终答案了。
Action: Final Answer
Action Input: 最终答案

历史对话和最新任务:
{messages} 
{agent_scratchpad}
"""

# 从原始的 Hub Prompt 切换到手动定义的 ChatPromptTemplate
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", prompt_template),
        # 注意：这里我们使用 'messages' 占位符来接收整个消息历史
    ]
).partial(tool_names=", ".join([t.name for t in tools])) # 预先填充 tool_names


# ----------------------------------------------------------
# 二、模型定义
# ----------------------------------------------------------
#  讯飞星火 "xdeepseekv32exp"
model_name = "xdeepseekv32exp"
# 通义
#model_name = "Qwen/Qwen2.5-VL-7B-Instruct"

# 定义大语言模型
llm = ChatOpenAI(model=model_name)

# 创建 ReAct 智能体执行器
# 会自动管理“思考 → 行动 → 观察 → 再思考”的循环逻辑
agent_executor = create_react_agent(llm, tools, prompt=prompt)


# ----------------------------------------------------------
# 三、定义状态结构（State）
# ----------------------------------------------------------
# LangGraph 中每个节点执行后会更新 State（类似状态机的上下文）
# PlanExecute 用于在节点间传递任务信息
class PlanExecute(BaseModel):
    input: str = Field(description="用户录入提示词")

    # 修正 1：设置默认值，让 Pydantic 允许其缺失
    plan: List[str] = Field(
        default_factory=list,
        description="用户录入提示词如果是大问题，会被拆分成几个小问题"
    )

    # 修正 2：设置默认值，让 Pydantic 允许其缺失
    past_steps: Annotated[List[Tuple], operator.add] = Field(
        default_factory=list,
        description="已执行的步骤及结果"
    )

    # 修正 3：设置默认值
    response: str = Field(
        default="",
        description="已执行的步骤及结果"
    )

# ----------------------------------------------------------
# 四、定义 Pydantic 数据模型
# ----------------------------------------------------------
# 定义一个计划模型类，用于描述未来要执行的计划
class Plan(BaseModel):
    """表示一组待执行的计划步骤"""
    steps: List[str] = Field(description="需要执行的不同步骤，应该按顺序排列")


# 定义 JSON 解析器
planner_parser = JsonOutputParser(pydantic_object=Plan)

# 定义生成计划的提示模板（系统角色）
planner_prompt = ChatPromptTemplate.from_messages([
    ("system",
     """
       对于给定的目标，提出一个简单的逐步计划。
       - 每个步骤应是独立、可执行的任务；
       - 不要添加多余步骤；
       - 最后一步应能直接得到最终答案；
       - 确保每步信息足够，不要跳过。
       
        你必须严格按照以下的 JSON Schema 格式返回结果：
       {format_instructions}
       """),
    ("placeholder", "{messages}")  # 这里插入用户输入
])

# 构建 planner（计划生成器）
# 注意：用 .partial() 注入格式指令，并使用 | planner_parser 进行解析
planner = planner_prompt.partial(format_instructions=planner_parser.get_format_instructions()) | ChatOpenAI(model=model_name, temperature=0) | planner_parser


# 移除原有使用 .with_structured_output 的代码
# planner = planner_prompt | ChatOpenAI(model=model_name, temperature=0).with_structured_output(Plan)


# 定义最终响应模型
class Response(BaseModel):
    """表示模型的直接回答"""
    response: str


# 定义 Act 模型，表示要执行的行为
# BaseModel 是 Pydantic 的基类，继承它后，类中的每个字段都会自动生成类型校验，序列号能力
# Union 是 Python 的类型注解，表示这个字段的值可以是Union中任意类型之一。即它既可以是直接回应（Response），也可以是新的计划（Plan）
# Field(...) 是 Pydantic 提供的函数，用来给字段添加元数据（例如描述、默认值、约束条件）,Field 会作为输出约束给大模型（langchain会将它转换为提示词）
class Act(BaseModel):
    """要执行的行为"""
    # action: Union[Response, Plan] = Field(description="要执行的行为。如果要回应用户，使用Response。 如果需要进一步使用工具获取答案，使用Plan。")
    action: Union[Response, Plan] = Field(description="若已可回答则返回Response，否则返回新的Plan。")


# 定义 JSON 解析器
replanner_parser = JsonOutputParser(pydantic_object=Act)
# 定义“重新规划”的提示模板 (在模板中加入 JSON 格式要求)
replanner_prompt = ChatPromptTemplate.from_template(
    """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。

你的目标是：
{input}

你的原计划是：
{plan}

你目前已完成的步骤是：
{past_steps}

相应地更新你的计划。如果不需要更多步骤并且可以返回给用户，那么就这样回应。如果需要，填写计划。只添加仍然需要完成的步骤。不要返回已完成的步骤作为计划的一部分。

你必须严格按照以下的 JSON Schema 格式返回结果：
{format_instructions}
"""
)

# 构建 replanner（重新规划器）
replanner = replanner_prompt.partial(format_instructions=replanner_parser.get_format_instructions()) | ChatOpenAI(model=model_name) | replanner_parser


# 移除原有使用 .with_structured_output 的代码
# replanner = replanner_prompt | ChatOpenAI(model=model_name).with_structured_output(Act)

# ----------------------------------------------------------
# 五、定义工作流节点函数
# ----------------------------------------------------------
async def main():
    # ---------- 1️⃣ 计划节点 ----------
    async def plan_step(state: PlanExecute):
        """根据用户输入生成计划"""
        plan = await planner.ainvoke({"messages": [("user", state.input)]})
        print(f"生成的计划: {plan}")
        return {"plan": plan["steps"]}

    # ---------- 2️⃣ 执行节点 ----------
    async def execute_step(state: PlanExecute):
        """执行计划中的第一个步骤"""

        plan = state.plan
        past_steps = state.past_steps  # 从状态中获取历史步骤

        # 1. 构造 Agent Scratchpad (历史记录)
        # LangChain 的 ReAct Agent 期望历史记录是 AIMessage 和 ToolMessage 的列表
        agent_scratchpad = []

        # 遍历已完成的 (任务, 结果) 对，将其转换为 ReAct 消息历史
        for task, response in past_steps:
            # 假设 Agent 思考后，会使用 ToolCall（我们不知道具体的Action/ToolName）
            # 我们可以简化为： Agent (AIMessage) -> Tool (ToolMessage)

            # 简化处理：将历史步骤作为通用的 AIMessage 和 ToolMessage 对
            # 这是为了模拟 ReAct 循环中的 "Thought / Action" (AIMessage) 和 "Observation" (ToolMessage)

            # 思考和行动（简化为 Agent 的结果，因为我们没有存储原始Thought/Action）
            agent_scratchpad.append(AIMessage(content=f"Thought: 我已完成任务 '{task}'.\nAction: Final Answer\nAction Input: {response}"))

            # 观察结果 (可选，但为了完整性)
            # agent_scratchpad.append(ToolMessage(content=response, tool_call_id="mock_id"))

        # 2. 构造当前任务
        plan_str = "\n".join(f"{i + 1}.{step}" for i, step in enumerate(plan))
        task = plan[0]
        task_formatted = f"""
    以下是计划：
    {plan_str}\n\n 你的任务是执行第1步任务，{task}"""

        # 3. 构造 Agent Executor 的完整输入
        # LangGraph Prebuilt Agent 期望将任务作为 HumanMessage 列表的一部分
        model_input = {
            "messages": [
                HumanMessage(content=task_formatted)
            ],
            # 💥 关键：传入构造好的 agent_scratchpad
            "agent_scratchpad": agent_scratchpad
        }

        agent_response = await agent_executor.ainvoke(model_input)

        # 4. 更新状态
        # 我们将任务和 Agent 最终的回复添加到 past_steps
        return {
            "past_steps": state.past_steps + [(task, agent_response["messages"][-1].content)]
        }

    # ---------- 3️⃣ 重新规划节点 ----------
    async def replan_step(state: PlanExecute):
        """根据已执行的步骤决定是否继续"""

        # output 是一个字典，例如：{'action': {'response': '...'}}
        output = await replanner.ainvoke(state)

        # 提取 'action' 字典
        action_data = output.get("action", {})

        # 判断 action_data 中是否包含 'response' 键 (Response 模型)
        if "response" in action_data:
            # 返回最终答案
            return {"response": action_data["response"]}

        # 判断 action_data 中是否包含 'steps' 键 (Plan 模型)
        elif "steps" in action_data:
            # 返回新的计划步骤
            return {"plan": action_data["steps"]}

        else:
            # 如果模型返回了不符合 Plan 或 Response 结构的 JSON，则抛出错误
            raise ValueError(f"Replanner 返回的 Act 结构不明确或不正确: {output}")

    # ---------- 4️⃣ 条件判断 ----------
    def should_end(state: PlanExecute) -> Literal["agent", "__end__"]:
        """判断是否结束"""
        if "response" in state and state.response:
            return "__end__"
        else:
            return "agent"

    # ----------------------------------------------------------
    # 六、构建 LangGraph 工作流
    # ----------------------------------------------------------
    # 定义状态图，PlanExecute 为全局共享状态结构
    workflow = StateGraph(PlanExecute)

    # 添加节点（即每个逻辑阶段）

    workflow.add_node("planner", plan_step)  # 添加计划节点
    workflow.add_node("agent", execute_step)  # 添加执行步骤节点
    workflow.add_node("replan", replan_step)  # 添加重新计划节点

    # 定义节点间的有向边（执行路径）
    workflow.add_edge(START, "planner")  # 开始 → 计划
    workflow.add_edge("planner", "agent")  # 计划 → 执行
    workflow.add_edge("agent", "replan")  # 执行 → 重新规划

    # 条件边：判断是否继续还是结束
    workflow.add_conditional_edges("replan", should_end)

    # 编译生成可运行的 LangGraph 应用
    app = workflow.compile()

    # 可视化输出状态图结构（Mermaid 格式）
    graph_png = app.get_graph().draw_mermaid_png()
    with open("llm_langgraph_1_agent工作流.png", "wb") as f:
        f.write(graph_png)

    # ----------------------------------------------------------
    # 七、运行工作流
    # ----------------------------------------------------------
    config = {"recursion_limit": 50}  # 设置递归上限防止死循环
    # 输入参数
    inputs = {"input": "2024年巴黎奥运会100米自由泳决赛冠军的家乡是哪里？请用中文答复"}

    # 异步流式运行工作流（可实时看到状态变化）
    async for event in app.astream(inputs, config=config):
        for k, v in event.items():
            if k != "__end__":
                print(v)



async def test_planner_output():
    """异步函数：用于测试 planner 链的输出类型和结构"""

    print("--- 正在测试 planner 链的异步调用 ---")

    # 异步调用 planner.ainvoke()，并使用 await 等待结果
    result_plan = await planner.ainvoke({"messages": [("user", "2024年巴黎奥运会100米自由泳决赛冠军的家乡是哪里？请用中文答复")]})

    # 打印返回结果的类型
    print(f"返回结果类型: {type(result_plan)}")

    # 打印返回结果的内容，检查其结构是否符合预期 (例如是否包含 'steps' 键)
    print(f"返回结果内容: {result_plan}")

    # 确保它是一个字典，并且包含 'steps' 键
    if isinstance(result_plan, dict) and 'steps' in result_plan:
        print("\n✅ 测试成功：返回了预期的字典结构，可用于 LangGraph 状态合并。")
    else:
        print("\n❌ 测试失败：返回结果不是预期的字典或缺少 'steps' 键。")


#使用 asyncio.run() 运行异步测试函数
#注意：asyncio.run() 接受一个协程对象，而不是 lambda 表达式
#if __name__ == "__main__":
#   asyncio.run(test_planner_output())


# 运行异步函数
asyncio.run(main())